# 这个文件用来把brpc的结果放到brpc_transfer_rates.csv里面。

import subprocess
import os
import csv

# 创建结果目录
results_dir = '../results_raw'
os.makedirs(results_dir, exist_ok=True)

# 结果文件
output_file = os.path.join(results_dir, 'brpc_transfer_rates.csv')

# 定义不同的参数组合
brpc_task_bytes = [2**i for i in range(3, 22)]  # 从8到2^20的所有2的幂次方

def convert_to_us(value, unit):
    """将时间单位转换为微秒"""
    if unit == 's':
        return value * 1_000_000
    elif unit == 'ms':
        return value * 1_000
    elif unit == 'us':
        return value
    else:
        raise ValueError(f"Unknown unit: {unit}")

def calculate_speed(total_size_bytes, time_us):
    """计算传输速度，单位为M/s"""
    if time_us == 0:
        return float('inf')  # 防止除以零
    speed_mbps = (total_size_bytes / time_us) * 1_000_000 / (1024 * 1024)
    return speed_mbps

# 打开文件以写入结果
with open(output_file, 'w', newline='') as csvfile:
    fieldnames = ['brpc_task_bytes', 'P50_us', 'P99_us', 'Speed_Mbps']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

    writer.writeheader()

    for brpc_task_byte in brpc_task_bytes:
        # 提示当前运行的参数组合
        print(f"Running with brpc_task_bytes={brpc_task_byte}")

        # 构建命令
        command = [
            './../../bin_dpu/test_communication_client',
            f'--brpc_byte={brpc_task_byte}'
        ]

        # 执行命令并捕获输出
        try:
            result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding='utf-8', errors='ignore')
        except Exception as e:
            print(f"Error executing command: {e}")
            continue

        # 解析输出
        stdout = result.stdout
        lines = stdout.splitlines()
        for line in lines:
            if 'brpc_byte' in line:
                parts = line.split('|')
                if len(parts) >= 4:
                    # 提取p50和p99值，并统一单位为us
                    p50_value, p50_unit = parts[2].strip().split()
                    p99_value, p99_unit = parts[3].strip().split()

                    p50_us = convert_to_us(float(p50_value), p50_unit)
                    p99_us = convert_to_us(float(p99_value), p99_unit)

                    # 计算实际传输大小
                    total_size_bytes = brpc_task_byte

                    # 计算速度
                    speed_mbps = calculate_speed(total_size_bytes, p50_us)

                    writer.writerow({
                        'brpc_task_bytes': brpc_task_byte,
                        'P50_us': p50_us,
                        'P99_us': p99_us,
                        'Speed_Mbps': speed_mbps
                    })

print(f"Results saved to {output_file}")
